# 导入必要的库
import pandas as pd
import numpy as np
import yfinance as yf # 导入yfinance库
from sklearn.linear_model import LinearRegression # 导入线性回归模型

# 使用yfinance获取数据，假设股票代码为NVDA，时间范围为2022年1月1日至今
# stock=input("请输入股票编号：")
stock="601669.SS"
df = yf.download(stock, period="MAX")
print(df.head())
# 定义你的持仓数量和成本价
position = int(input("请输入持仓：")) # 持仓数量
# cost = input("请输入成本价：") # 成本价
cost = float(input("请输入成本：")) # 成本价
# 计算5日均线和20日均线
df["MA5"] = df["Close"].rolling(5).mean()
df["MA20"] = df["Close"].rolling(20).mean()

# 定义买卖信号，当5日均线上穿20日均线时为买入信号，当5日均线下穿20日均线时为卖出信号
df["signal"] = np.where(df["MA5"] > df["MA20"], 1, -1)
df["signal"] = df["signal"].diff().fillna(0)

# 计算策略收益率，假设每次买入或卖出100股，不考虑交易成本和滑点
df["return"] = df["Close"].pct_change().fillna(0)
df["strategy_return"] = df["return"] * df["signal"].shift(1) * 100



# 计算你的持仓市值和盈亏
df["position_value"] = position * df["Close"]
df["profit_loss"] = df["position_value"] - position * cost

# 使用线性回归模型来预测明天的收盘价
X = np.arange(len(df)).reshape(-1, 1) # 特征变量为日期序号
y = df["Close"].values # 目标变量为收盘价
model = LinearRegression() # 创建线性回归模型对象
model.fit(X, y) # 训练模型
X_pred = np.array([len(df)]).reshape(-1, 1) # 预测变量为明天的日期序号
y_pred = model.predict(X_pred) # 预测明天的收盘价

# 计算明天的5日均线和20日均线
ma5_pred = (df["Close"].tail(4).sum() + y_pred) / 5 # 明天的5日均线等于最近四天的收盘价加上明天的收盘价除以5
ma20_pred = (df["Close"].tail(19).sum() + y_pred) / 20 # 明天的20日均线等于最近十九天的收盘价加上明天的收盘价除以20

# 判断明天的买卖信号，如果5日均线上穿20日均线，则为买入信号；如果5日均线下穿20日均线，则为卖出信号；否则无信号
if ma5_pred > ma20_pred and df["MA5"].iloc[-1] <= df["MA20"].iloc[-1]:
    signal_pred = "buy"
elif ma5_pred < ma20_pred and df["MA5"].iloc[-1] >= df["MA20"].iloc[-1]:
    signal_pred = "sell"
else:
    signal_pred = "hold"

# 打印预测结果，并翻译成中文
print(f"Tomorrow's predicted close price is {y_pred[0]:.2f}") #

# 根据你的持仓和成本，明天是否进行买卖，买卖价格多少，并翻译成中文
if signal_pred == "buy":
    print(f"You should buy 100 shares at {y_pred[0]:.2f}") # 你应该以预测的收盘价买入100股
    print(f"The predicted buy price is {y_pred[0]:.2f}") # 预测的买入价是
elif signal_pred == "sell":
    print(f"You should sell 100 shares at {y_pred[0]:.2f}") # 你应该以预测的收盘价卖出100股
    print(f"The predicted sell price is {y_pred[0]:.2f}") # 预测的卖出价是
else:
    print(f"You should hold your position") # 你应该保持你的持仓